import torch
import torch.nn as nn
import numpy as np
import pandas as pd
import logging
import dgl
import os
import json
import functools
import traceback
from tqdm import tqdm
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Tuple, Set, Optional, Any, Union
from collections import defaultdict
import shutil

# 导入项目相关模块
from source.data_processor.utils.experiment_data import ExperimentData
from source.localize.graph.systemDependencyGraph import SystemDependencyGraphBuilder
from .system_feature import SystemFeatureVector

# 导入重构后的辅助模块
from .utils.data_processor import DataProcessor
from .utils.graph_builder import GraphProcessor
from .utils.model_trainer import ModelTrainer
from .utils.feature_extractor import FeatureExtractor
from .utils.config_manager import ConfigManager

# 配置日志
logging.basicConfig(
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    level=logging.INFO,
)
logger = logging.getLogger(__name__)


class ARTMiner:
    """
    ART模型训练和特征提取器

    该类将系统依赖关系和指标数据转换为机器学习模型的输入，训练模型，并提取服务级别偏差(SLD)特征。

    处理流程:
    1. 数据准备阶段 - 提取并预处理服务依赖关系和指标数据
    2. 模型训练阶段 - 使用准备好的数据训练ART模型
    3. 特征提取阶段 - 使用训练好的模型提取SLD和ILD特征
    4. 分析诊断阶段 - 根据特征计算相似度，识别潜在的异常组件
    """

    def __init__(
        self,
        config: Dict[str, Any] = None,
        output_dir: Optional[Union[str, Path]] = None,
    ):
        """
        初始化ARTMiner

        参数:
            config: 配置字典，包含模型参数和训练参数
            output_dir: 输出目录，用于保存模型和特征
        """
        self.config = config or {}
        self.output_dir = Path(output_dir) if output_dir else None

        # 创建输出目录
        if self.output_dir:
            self._clear_directory(self.output_dir)
            self.output_dir.mkdir(parents=True, exist_ok=True)
            self.model_dir = self.output_dir / "models"
            self.feature_dir = self.output_dir / "features"
            self.model_dir.mkdir(exist_ok=True)
            self.feature_dir.mkdir(exist_ok=True)

        # 初始化子模块
        self.data_processor = DataProcessor(
            window_size=self.config.get("window_size", 6),
            max_gap=self.config.get("max_gap", 60),
        )
        self.graph_processor = GraphProcessor()
        self.model_trainer = ModelTrainer(config)
        self.feature_extractor = FeatureExtractor(self.output_dir)
        self.config_manager = ConfigManager(self.output_dir)

        # 内部状态
        self.model = None
        self.graph = None
        self.train_samples = []
        self.test_samples = []
        self.cmdb_id_to_idx = {}
        self.idx_to_cmdb_id = {}
        self.useful_entities = set()
        self.metric_to_feature_idx = {}
        self.sld_features = None
        self.experiment_data = None
        self.feature_vector = None

        # 数据处理配置
        self.data_config = {
            "window_size": self.config.get("window_size", 6),
            "max_gap": self.config.get("max_gap", 60),
        }

    def _clear_directory(self, directory: Path) -> None:
        """
        清空指定目录中的所有内容
        参数:
            directory: 要清空的目录路径
        """
        if not directory.exists():
            return
        logger.info(f"清空目录: {directory}")
        # 删除目录中的所有文件和子目录
        for item in directory.iterdir():
            if item.is_dir():
                shutil.rmtree(item)
            else:
                item.unlink()

    def prepare_data(
        self,
        experiment_data: ExperimentData,
        graph_builder: SystemDependencyGraphBuilder,
    ) -> bool:
        """
        准备训练和测试数据

        数据准备阶段的详细步骤:
        1. 提取有用实体和服务依赖关系
        2. 创建指标映射和构建网络图
        3. 准备节点映射和调整模型参数
        4. 分割数据并转换为模型输入样本

        参数:
            experiment_data: ExperimentData对象
            graph_builder: SystemDependencyGraphBuilder对象

        返回:
            bool: 准备成功返回True，否则返回False
        """
        logger.info("开始准备ART模型的训练和测试数据")
        self.experiment_data = experiment_data

        # 步骤1: 提取有用实体和服务依赖关系
        logger.info("第1步: 检查图构建器和提取接口依赖关系")
        graph_id = graph_builder.active_graph_id

        # 获取接口依赖关系
        interface_deps = graph_builder.get_interface_dependencies() or {}

        # 提取有用实体 - 从experiment_data中提取而不是从graph_builder中提取
        self.useful_entities = self.graph_processor.extract_useful_entities(
            experiment_data
        )

        # 步骤2: 创建指标映射和构建网络图
        logger.info("第2步: 获取和预处理指标数据")
        metrics_df = experiment_data.metrics_df

        # 确保timestamp为datetime类型
        metrics_df = self.data_processor.ensure_datetime(metrics_df)

        # 创建指标到特征索引的映射
        self.metric_to_feature_idx = self.data_processor.create_metric_mapping(
            metrics_df, self.useful_entities
        )
        if not self.metric_to_feature_idx:
            logger.error("未能创建有效的指标映射")
            return False

        # 构建网络图和节点映射
        self.graph, all_cmdb_ids, self.cmdb_id_to_idx, self.idx_to_cmdb_id = (
            self.graph_processor.prepare_graph(
                interface_deps, graph_builder, self.useful_entities
            )
        )

        # 步骤3: 更新模型参数
        logger.info("第3步: 更新模型参数")
        self.model_trainer.update_model_config(
            channel_dim=len(self.metric_to_feature_idx), instance_dim=len(all_cmdb_ids)
        )

        # 步骤4: 分割数据并转换为模型输入样本
        logger.info("第4步: 分割异常前后的数据并转换为训练样本")

        # 确定滑动窗口大小，可以根据实际情况调整
        window_size = self.config.get("window_size", 10)

        pre_anomaly_data, post_anomaly_data = self.data_processor.split_data_by_anomaly(
            experiment_data, window_size
        )

        # 转换训练样本
        self.train_samples = self.data_processor.convert_metrics_to_samples(
            pre_anomaly_data.metrics_df,
            self.graph,
            self.cmdb_id_to_idx,
            self.metric_to_feature_idx,
            self.useful_entities,
        )
        if not self.train_samples:
            logger.error("未能生成训练样本")
            return False

        # 转换测试样本
        self.test_samples = self.data_processor.convert_metrics_to_samples(
            post_anomaly_data.metrics_df,
            self.graph,
            self.cmdb_id_to_idx,
            self.metric_to_feature_idx,
            self.useful_entities,
        )
        if not self.test_samples:
            logger.warning("未能生成测试样本，将使用训练样本的一部分作为测试样本")
            test_size = max(1, len(self.train_samples) // 5)
            self.test_samples = self.train_samples[-test_size:]

        logger.info(
            f"数据准备完成: {len(self.train_samples)}个训练样本, {len(self.test_samples)}个测试样本"
        )
        return True

    def train_model(self) -> bool:
        """
        使用准备好的训练数据训练模型

        训练步骤:
        1. 验证训练样本是否准备好
        2. 训练ART模型
        3. 保存训练好的模型

        返回:
            bool: 训练成功返回True，否则返回False
        """
        # 验证训练样本
        if not self.train_samples:
            logger.error("训练样本为空，无法训练模型")
            return False

        logger.info("开始训练ART模型")

        # 训练模型
        self.model = self.model_trainer.train_model(self.train_samples)
        if self.model is None:
            logger.error("模型训练失败")
            return False

        # 保存模型
        if self.output_dir and self.experiment_data:
            anomaly_id = (
                self.experiment_data.anomaly_id if self.experiment_data else None
            )
            self.model_trainer.save_model(self.model, self.model_dir, anomaly_id)

        return True

    def extract_unified_features(
        self, top_k_value: int = 3, auto_save: bool = False
    ) -> SystemFeatureVector | None:
        """
        使用新的统一特征提取方法，同时提取SLD和ILD特征

        参数:
            top_k_value: 选择的top-k实例数量
            auto_save: 是否自动保存特征和结果

        返回:
            SystemFeatureVector | None: 特征向量对象，如果提取失败则返回None
        """
        if not self.model or not self.test_samples:
            logger.error("模型或测试样本未准备好，无法提取特征")
            return None

        return self.feature_extractor.extract_sld_ild_features(
            model=self.model,
            test_samples=self.test_samples,
            idx_to_cmdb_id=self.idx_to_cmdb_id,
            top_k_value=top_k_value,
            anomaly_id=(
                self.experiment_data.anomaly_id if self.experiment_data else None
            ),
            output_dir=str(self.feature_dir) if self.feature_dir else None,
            auto_save=auto_save,
        )

    def analyze(
        self,
        experiment_data: ExperimentData,
        graph_builder: SystemDependencyGraphBuilder,
        top_k_value: int = 3,
        save_results: bool = True,
    ) -> Dict[str, Any]:
        """
        运行完整的特征提取流程，使用新的统一特征提取方法

        参数:
            experiment_data: 实验数据对象
            graph_builder: 系统依赖图构建器对象
            top_k_value: 选择的top-k实例数量
            save_results: 是否保存结果

        返回:
            Dict[str, Any]: 包含分析结果的字典
        """
        start_time = datetime.now()
        logger.info(f"开始特征提取流程，top_k={top_k_value}")

        # 第一步: 数据准备
        if not self.prepare_data(experiment_data, graph_builder):
            return {"success": False, "error": "数据准备失败"}

        # 第二步: 模型训练
        if not self.train_model():
            return {"success": False, "error": "模型训练失败"}

        # 第三步: 特征提取
        result = self.extract_unified_features(
            top_k_value=top_k_value, auto_save=save_results
        )

        if result is None:
            return {"success": False, "error": "特征提取失败"}

        # 构建结果对象
        analysis_result = {
            "success": True,
            "runtime": (datetime.now() - start_time).total_seconds(),
        }

        # 保存结果
        self.feature_vector = result
        analysis_result.update(
            {
                "feature_vector": result,
                "sld_tensor": result.sld_tensor,
                "ild_tensor": result.ild_tensor,
                "similarity_results": result.similarity_results,
            }
        )

        logger.info(f"特征提取流程完成，用时 {analysis_result['runtime']:.2f} 秒")

        return analysis_result
